import streamlit as st
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from code import code_five

# 页面设置
st.set_page_config(layout="wide")
st.title("⚡ 优化器 (Optimizers)")

# ========== 侧边栏选择 ==========
opt_choice = st.sidebar.selectbox(
    "选择优化器",
    ["GD", "SGD", "MBGD", "Momentum", "Nesterov", "AdaGrad", "RMSProp", "Adam", "总结"]
)

# ========== 优化器介绍 ==========
descriptions = {
    "GD": {
        "原理": "批量梯度下降，每次使用全部样本计算梯度。",
        "公式": "$ \\theta = \\theta - \\eta \\, \\nabla J(\\theta) $",
        "优点": "方向稳定。",
        "缺点": "计算开销大，不适合大数据。"
    },
    "SGD": {
        "原理": "每次使用单一样本计算梯度，更新参数。",
        "公式": "$ \\theta = \\theta - \\eta \\, \\nabla J_i(\\theta) $",
        "优点": "实现简单，适合大数据。",
        "缺点": "收敛不稳定。"
    },
    "MBGD": {
        "原理": "小批量梯度下降（Mini-Batch GD），每次使用一小批样本计算梯度，是 GD 和 SGD 的折中。",
        "公式": "$ \\theta = \\theta - \\eta \\, \\frac{1}{m} \\sum_{i=1}^m \\nabla J_i(\\theta) $",
        "优点": "兼顾收敛稳定性和计算效率，是实际应用中最常用的方法。",
        "缺点": "批量大小需要调节，不同任务最优批量大小不同。"
    },
    "Momentum": {
        "原理": "引入动量项，累积历史梯度。",
        "公式": "$ v_t = \\mu v_{t-1} - \\eta g_t, \\\\; \\theta = \\theta + v_t $",
        "优点": "加速收敛，减小震荡。",
        "缺点": "需调节动量系数。"
    },
    "Nesterov": {
        "原理": "在 Momentum 基础上提前预估梯度方向。",
        "公式": "$ v_t = \\mu v_{t-1} - \\eta \\, \\nabla J(\\theta + \\mu v_{t-1}) $",
        "优点": "收敛速度更快。",
        "缺点": "实现稍复杂。"
    },
    "AdaGrad": {
        "原理": "对不同参数使用不同的自适应学习率。",
        "公式": "$ \\theta = \\theta - \\frac{\\eta}{\\sqrt{G_t}+\\epsilon} g_t $",
        "优点": "适合稀疏特征。",
        "缺点": "学习率单调下降过快。"
    },
    "RMSProp": {
        "原理": "改进 AdaGrad，使用指数加权平均梯度平方。",
        "公式": "$ E[g^2]_t = \\gamma E[g^2]_{t-1} + (1-\\gamma)g_t^2 $",
        "优点": "解决学习率过早衰减。",
        "缺点": "需调节衰减率 γ。"
    },
    "Adam": {
        "原理": "结合 Momentum 和 RMSProp。",
        "公式": r"""
        $ m_t = \beta_1 m_{t-1} + (1-\beta_1) g_t $  
        $ v_t = \beta_2 v_{t-1} + (1-\beta_2) g_t^2 $  
        $ \hat{m}_t = m_t / (1-\beta_1^t), \hat{v}_t = v_t / (1-\beta_2^t) $  
        $ \theta = \theta - \eta \hat{m}_t / (\sqrt{\hat{v}_t}+\epsilon) $
        """,
        "优点": "收敛快，效果稳定。",
        "缺点": "对超参数较敏感。"
    }
}

# ========== 如果选择总结 ==========
if opt_choice == "总结":
    st.subheader("📊 优化器对比总结")
    data = {
        "优化器": ["GD", "SGD", "MBGD", "Momentum", "AdaGrad", "RMSProp", "Adam"],
        "核心思想": [
            "使用全部数据计算梯度",
            "使用单个样本更新",
            "使用小批量样本",
            "引入动量项",
            "自适应学习率",
            "指数加权平均",
            "结合Momentum和RMSProp"
        ],
        "优点": [
            "收敛稳定",
            "计算快，可在线学习",
            "平衡速度与稳定性",
            "加速收敛，减少震荡",
            "稀疏数据表现好（NLP）",
            "解决AdaGrad学习率问题",
            "自适应学习率，收敛快"
        ],
        "缺点": [
            "计算慢，内存要求高",
            "波动大，收敛不稳定",
            "需要调整批量大小",
            "需要调整动量参数",
            "学习率持续下降",
            "需要调整衰减率",
            "超参数较多"
        ],
        "适用场景": [
            "小数据集",
            "大规模数据",
            "大多数深度学习任务",
            "高维、稀疏梯度",
            "自然语言处理",
            "非平稳目标函数",
            "大多数深度学习应用"
        ]
    }
    st.table(pd.DataFrame(data))

    st.markdown("---")
    st.subheader("📝 优化器选择建议")
    st.markdown("""
    1. **Adam**：默认推荐选择，在大多数情况下表现良好  
    2. **SGD**：配合学习率调度器，在许多任务中可获得更好的最终性能  
    3. **AdaGrad**：在稀疏数据场景（如 NLP）中表现优异  
    4. **GD / MBGD**：对于简单任务或小数据集即可使用  
    """)

    st.stop()


with st.form("optimizer_form"):
    # ========== 展示单个优化器 ==========
    info = descriptions[opt_choice]
    st.markdown(f"## 📘 {opt_choice} 优化器")
    st.markdown(f"**原理**：{info['原理']}")
    st.markdown(f"**公式**：{info['公式']}")
    st.markdown(f"**优点**：{info['优点']}")
    st.markdown(f"**缺点**：{info['缺点']}")
    st.markdown("---")

    # ========== 代码示例 ==========
    st.subheader("💻 PyTorch 代码示例")
    sample_code = code_five.optimizer_code(opt_choice)
    st.code(sample_code, language="python")

    st.markdown("---")

    # ========== 等高线梯度下降演示 ==========
    st.subheader("📉 等高线下的优化路径演示")
    code_five.optimizer(opt_choice)
    col1, col2 = st.columns([10, 1])
    with col2:
        submit = st.form_submit_button("我已学习")
    if submit:
        st.success("恭喜你，已掌握该优化器的使用方法！")
        st.balloons()
